package priv.tuyou.redis.limit;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;

/**
 * 漏斗法-限流器
 *
 * 使用一个大小为limit的队列实现漏斗，value存放入队时间，rpop弹出入队超过一个周期的元素（惰性移除），如果剩余空间足够则lpush入队
 * @author tuyou_cd
 * @date 2021/4/23 15:47
 */
public class FunnelRateLimiter extends AbstractRateLimiter {

    public FunnelRateLimiter(JedisPool pool, int period, int limit) {
        super(pool, period, limit);
    }

    @Override
    public boolean visit(int count) {
        removeExpired();
        try (Jedis jedis = pool.getResource()) {
            if (jedis.llen(key) >= limit) {
                return false;
            }

            jedis.lpush(key, String.valueOf(System.currentTimeMillis()));
        }
        return true;
    }

    private void removeExpired() {
        try (Jedis jedis = pool.getResource()) {
            while (true) {
                String inQueueTime = jedis.lindex(key, -1);
                if (inQueueTime == null) {
                    break;
                }

                if (System.currentTimeMillis() - Long.parseLong(inQueueTime) >= period * 1000) {
                    jedis.rpop(key);
                } else {
                    break;
                }
            }
        }
    }
}
